from faker import Faker
import random
from datetime import datetime, timedelta
from dateutil.relativedelta import relativedelta
from data_mock.utils import FileUtil, MysqlUtils_AI

fake = Faker('zh_CN')

# ========== 基础数据配置 ==========
CANCER_DIAGNOSES = [
    ("C34.901", "肺恶性肿瘤", "T2N1M0", "主要诊断"),
    ("C50.902", "乳腺恶性肿瘤", "T1N0M0", "主要诊断"),
    ("C16.901", "胃恶性肿瘤", "T3N2M0", "主要诊断"),
    ("C18.901", "结肠恶性肿瘤", "T4N1M1", "主要诊断"),
    ("D63.1*", "肿瘤性贫血", None, "并发症诊断"),
    ("I10.x00", "高血压病", None, "其他诊断"),
    ("E11.900", "2型糖尿病", None, "其他诊断"),
    ("J18.900", "肺炎", None, "医院感染诊断")
]

DIAG_TYPES = ["入院诊断", "出院诊断", "病理诊断", "并发症诊断", "医院感染诊断"]
HOSPITAL_INFO = [
    ("H1001", "北京协和医院", "P01", "东院"),
    ("H1002", "上海瑞金医院", "P02", "总院"),
    ("H1003", "广州中山医院", "P03", "肿瘤分院")
]


# ========== 时间范围配置 ==========
def generate_day_range(start_day, end_day):
    start = datetime.strptime(start_day, "%Y-%m-%d")
    end = datetime.strptime(end_day, "%Y-%m-%d")
    return [
        (start + timedelta(days=i)).strftime("%Y-%m-%d")
        for i in range((end - start).days + 1)
    ]


DAY_RANGE = generate_day_range("2025-08-01", "2025-08-03")


def generate_records_per_day(day: str):
    date_obj = datetime.strptime(day, '%Y-%m-%d')
    return 10 if date_obj.day % 2 == 0 else 6


# ========== 核心生成函数 ==========
def generate_admission_diagnosis_records():
    sql_statements = []
    diag_id_counter = 50001

    # 先获取已生成的入院记录（模拟关联）
    admission_records = [
        {"record_sn": f"REC{10001 + i}", "patient_id": f"PT{2001 + i}", "visit_sn": f"VIS{2001 + i}"}
        for i in range(50)
    ]

    for day in DAY_RANGE:
        RECORDS_PER_day = generate_records_per_day(day)
        for i in range(1, RECORDS_PER_day + 1):
            # 随机选择关联的入院记录
            admission = random.choice(admission_records)

            # 医院信息
            hospital_code, hospital_name, branch_code, branch_name = random.choice(HOSPITAL_INFO)

            # 每个患者生成1-5个诊断
            diag_count = random.randint(1, 5)
            for diag_seq in range(1, diag_count + 1):
                diag_id = f"DIAG{diag_id_counter}"
                diag_id_counter += 1

                # 第一个诊断必须为主要诊断
                if diag_seq == 1:
                    diag_code, diag_name, tnm_staging, diag_type = random.choice(
                        [d for d in CANCER_DIAGNOSES if d[3] == "主要诊断"]
                    )
                    maindiag_mark = "是"
                else:
                    diag_code, diag_name, _, diag_type = random.choice(CANCER_DIAGNOSES)
                    maindiag_mark = "否"

                # 时间信息
                record_datetime = fake.date_time_between(
                    start_date=datetime.strptime(day, "%Y-%m-%d"),
                    end_date=datetime.strptime(day, "%Y-%m-%d") + timedelta(days=1)
                ).strftime('%Y-%m-%d %H:%M:%S')

                # 生成yy_record_md5
                yy_record_md5 = f"md5_{random.getrandbits(128):032x}"

                data = {
                    'diag_code': diag_code,
                    'diag_id': diag_id,
                    'diag_name': diag_name,
                    'diag_seq': str(diag_seq),
                    'diag_type': diag_type,
                    'fk_record_sn': admission['record_sn'],
                    'from_table': 'EMR_DIAGNOSIS',
                    'from_yy_record_id': f"SOURCE_{random.randint(10000, 99999)}",
                    'hospital_code': hospital_code,
                    'hospital_name': hospital_name,
                    'maindiag_mark': maindiag_mark,
                    'patient_id': admission['patient_id'],
                    'patient_id_old': f"OLD_{admission['patient_id']}",
                    'record_datetime': record_datetime,
                    'record_status': 1,
                    'record_update_datetime': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                    'visit_sn': admission['visit_sn'],
                    'yy_collection_datetime': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                    'yy_record_md5': yy_record_md5,
                    'yy_upload_status': 0,
                    'yy_etl_time': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                    'yy_upload_time': None,
                    'yy_batch_time': day,
                    'yy_record_batch_id': f"BATCH{day}_{diag_id}",
                    'yy_backfill_time': None,
                    'yy_backfill_status': None,
                    'branch_code': branch_code,
                    'branch_name': branch_name,
                    'date_for_partition': record_datetime
                }

                sql = _generate_sql('b05_1_1', data)
                sql_statements.append(sql)

    return sql_statements


def _generate_sql(table, data):
    columns = ', '.join([f'`{k}`' for k in data.keys()])
    values = []
    for v in data.values():
        if v is None:
            values.append('NULL')
        elif isinstance(v, (int, float)):
            values.append(str(v))
        else:
            escaped_value = str(v).replace("'", "''")
            values.append(f"'{escaped_value}'")
    return f"INSERT INTO `{table}` ({columns}) VALUES ({', '.join(values)});"


# ========== 执行生成 ==========
if __name__ == "__main__":
    records = generate_admission_diagnosis_records()
    # 写入数据库
    MysqlUtils_AI.insert_data_to_hub(records, 'b05_1_1')
    # 或写入SQL文件
    # FileUtil.generate_sql_file(records, "b05_1_1_diagnosis_records.sql")